package com.muyu.danmo.util;

import lombok.RequiredArgsConstructor;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.springframework.data.domain.PageRequest;
import org.springframework.data.elasticsearch.core.ElasticsearchRestTemplate;
import org.springframework.data.elasticsearch.core.SearchScrollHits;
import org.springframework.data.elasticsearch.core.mapping.IndexCoordinates;
import org.springframework.data.elasticsearch.core.query.NativeSearchQuery;
import org.springframework.data.elasticsearch.core.query.NativeSearchQueryBuilder;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;
import java.util.function.Consumer;

/**
 * @author Administrator
 */
@Component
@RequiredArgsConstructor
public class ElasticScrollHelper {

    private final ElasticsearchRestTemplate elasticsearchRestTemplate;

    /**
     * 通用 scroll 查询工具
     *
     * @param indexName        索引名
     * @param scrollTimeoutMs  scroll 超时时间（毫秒）
     * @param pageSize         每页大小
     * @param clazz            返回类型
     * @param queryConsumer    条件构造函数
     * @param <T>              实体类
     * @return 查询到的全部结果
     */
    public <T> List<T> scrollQuery(
            String indexName,
            long scrollTimeoutMs,
            int pageSize,
            Class<T> clazz,
            Consumer<BoolQueryBuilder> queryConsumer
    ) {
        BoolQueryBuilder boolQuery = new BoolQueryBuilder();
        queryConsumer.accept(boolQuery);

        NativeSearchQuery searchQuery = new NativeSearchQueryBuilder()
                .withQuery(boolQuery)
                .withPageable(PageRequest.of(0, pageSize))
                .build();

        List<T> resultList = new ArrayList<>();

        SearchScrollHits<T> scrollHits = elasticsearchRestTemplate.searchScrollStart(
                scrollTimeoutMs,
                searchQuery,
                clazz,
                IndexCoordinates.of(indexName)
        );

        while (scrollHits != null && scrollHits.hasSearchHits()) {
            scrollHits.getSearchHits().forEach(hit -> resultList.add(hit.getContent()));
            scrollHits = elasticsearchRestTemplate.searchScrollContinue(
                    scrollHits.getScrollId(),
                    scrollTimeoutMs,
                    clazz,
                    IndexCoordinates.of(indexName)
            );
        }

        if (scrollHits != null) {
            elasticsearchRestTemplate.searchScrollClear(List.of(scrollHits.getScrollId()));
        }

        return resultList;
    }


    /**
     * 通用 count 查询工具
     *
     * @param indexName     索引名
     * @param clazz         实体类类型
     * @param queryConsumer 条件构造器
     * @param <T>           实体类
     * @return 匹配的文档数量
     */
    public <T> long countQuery(
            String indexName,
            Class<T> clazz,
            Consumer<BoolQueryBuilder> queryConsumer
    ) {
        BoolQueryBuilder boolQuery = new BoolQueryBuilder();
        queryConsumer.accept(boolQuery);

        NativeSearchQuery searchQuery = new NativeSearchQueryBuilder()
                .withQuery(boolQuery)
                .build();

        return elasticsearchRestTemplate.count(searchQuery, clazz, IndexCoordinates.of(indexName));
    }
}